from sqlalchemy import create_engine
import pandas as pd
import statsmodels.api as sm
import numpy as np
import matplotlib.pyplot as plt

# 数据库配置
db_config = {
    'host': '111.231.14.211',
    'user': 'tushare',
    'password': 'root',
    'database': 'tushare',
    'port': 13307,
    'charset': 'utf8mb4'
}

# 创建数据库引擎
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset=utf8mb4"
)

# SQL 查询
query = """
SELECT d.*, m.net_mf_vol , m.sell_elg_vol , m.buy_elg_vol, m.sell_lg_vol, m.buy_lg_vol , i.closes as i_closes, i.vol as i_vol
FROM date_1 d
JOIN moneyflows m ON d.ts_code = m.ts_code AND d.trade_date = m.trade_date
LEFT JOIN index_daily i ON d.trade_date = i.trade_date and i.ts_code='000001.SH'
WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' AND d.ts_code='000001.SZ'
"""

# 使用 SQLAlchemy 引擎执行查询并分块读取数据
chunk_size = 10000
df_chunks = pd.read_sql_query(query, engine, chunksize=chunk_size)
df1 = pd.concat(df_chunks, ignore_index=True)

# 计算 zd_close 和 hz_close 列
df1['zd_close'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2)
df1['hz_close'] = round((df1['i_closes'] - df1['i_closes'].shift(1)) / df1['i_closes'].shift(1), 2)
df1['hz_vol'] = round((df1['i_vol'].shift(1) - df1['i_vol'].shift(2)) / df1['i_vol'].shift(2), 2)

# 删除包含 NaN 值的行
df1 = df1.dropna(subset=['zd_close', 'hz_close', 'hz_vol']).reset_index(drop=True)

# 获取数值类型的列，并去除指定的列
numeric_cols = df1.select_dtypes(include=['number']).columns.tolist()
x = df1[['amount', 'net_mf_vol', 'sell_elg_vol', 'buy_lg_vol', 'hz_close', 'hz_vol']].copy()

# 确保没有 inf 或 NaN 值
x = x.replace([np.inf, -np.inf], np.nan).dropna()

# 执行 PCA 分析
eigenvalues, eigenvectors = np.linalg.eig(np.cov(x, rowvar=False))
print('累计贡献率为: ', round(eigenvalues[:5].sum() / eigenvalues.sum(), 4) * 100, '%')

n_components = 5
top_eigenvectors = eigenvectors[:, :n_components]

principal_components = np.dot(x, top_eigenvectors)
data_pca = pd.concat([df1, pd.DataFrame(principal_components, columns=[f'PC{i+1}' for i in range(n_components)])], axis=1)

X_pca = data_pca[[f'PC{i+1}' for i in range(n_components)]].copy()
X_pca = sm.add_constant(X_pca)

y = df1['zd_close'].copy()

model_pca = sm.OLS(y, X_pca)
result_pca = model_pca.fit()

print("\n回归模型结果:")
print(result_pca.summary())

# 选择部分主成分进行进一步分析
X_pca_selected = data_pca[['PC1', 'PC3', 'PC4']]
X_pca_selected = sm.add_constant(X_pca_selected)

model_pca_selected = sm.OLS(y, X_pca_selected)
result_pca_selected = model_pca_selected.fit()

print("\n回归模型结果:")
print(result_pca_selected.summary())

# 绘制散点图
fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))

for i, col in enumerate(X_pca_selected.columns[1:], start=1):
    axes[i-1].scatter(X_pca_selected[col], y, s=50, alpha=0.7)
    axes[i-1].set_xlabel(col)
    axes[i-1].set_ylabel('zd_close')
    axes[i-1].set_title(f'{col} vs zd_close')

plt.tight_layout()
plt.show()